import json
import os
from argparse import ArgumentParser

from tools.dataset_helper import open_img_np
from tools.search_feature_points import get_png_crn

def change_img2pt(img2pt_json: list, fn2wkt=None) -> None:
    """
    change img2pt file's keys from filename to key word
    """
    for j in img2pt_json:
        img2pt = _read_json(j)
        k = list(img2pt.keys())
        new_dict = {}
        for i in k:
            new_dict[fn2wkt(i)] = img2pt[i]

        with open(j, "w") as fp:
            json.dump(new_dict, fp, indent='\t', sort_keys=True)


def select_patch(*, img2pt_json: list, save_path: str, 
    pt_range: list, thresh={}, wkt2fn=None
) -> None:
    """
    select patch according to img2pt in which there are pts in all modals

    Param
    -----
    pt_range :  select patchs of which pts in range x~[l[0],l[1]) y~[l[2],l[3]) 
    """
    img2pt = []
    for json_ in img2pt_json:
        img2pt.append(_read_json(json_))
        
    length = len(img2pt)

    valid_img2pt = []
    # remove the pt of which response lower than thresh
    for dict_ in img2pt:
        if len(thresh) == 0:
            break
        for wkt in dict_:
            rm_list = []

            for k, v in dict_[wkt].items():
                r = v["response"]
                m = v["modal"]
                xy = v["xy"]
                if (m in thresh and r < thresh[m]) \
                or (
                    pt_range[0] > xy[0] or xy[0] >= pt_range[1] 
                or  pt_range[2] > xy[1] or xy[1] >= pt_range[3]
                ):
                    rm_list.append(k)

            for k in rm_list:
                dict_[wkt].pop(k)

        valid_img2pt.append(dict_)

    img2pt = valid_img2pt

    photo_set = set(filter(lambda a: len(img2pt[0][a])>0, img2pt[0].keys()))
    for i in range(1, length):
        tmp_set = set(filter(lambda a: len(img2pt[i][a])>0, img2pt[i].keys()))
        photo_set = photo_set.intersection(tmp_set)
        
        
    pt_dict = {}
    # 每张对应的图片
    for img in photo_set:
        pt_dict[img] = {}
        # 每个角点dict文件
        for dict_ in img2pt:
            pt_dict[img].update(dict_[img])

    result = {}
    result["OPT"] = []
    result["SAR"] = []
    for wkt, pts in pt_dict.items():
        # flag = False
        # modal_key = ""
        # for pt_attr in pts.values():
        #     modal_key = pt_attr["modal"]
        #     xy = pt_attr["xy"]
        #     if pt_range[0] <= xy[0] and xy[0] < pt_range[1] \
        #     and pt_range[2] <= xy[1] and xy[1] < pt_range[3]:
        #         flag = True
        #         break
        
        # if flag:
        #     path = wkt2fn(wkt, modal="opt") 
        #     result["OPT"].append(path)
        #     path = wkt2fn(wkt, modal="sar") 
        #     result["SAR"].append(path)

        path = wkt2fn(wkt, modal="opt") 
        result["OPT"].append(path)
        path = wkt2fn(wkt, modal="sar") 
        result["SAR"].append(path)

    with open(save_path, 'w') as fp:
        json.dump(result, fp, indent='\t', sort_keys=True)
    pass


def _read_json(json_dir):
    '''
    从json转化为dict

    Param:
    -----
    json_dir: 对应json路径

    Return:
    ----
    返回转化后的字典
    同时改变src_dict属性
    '''
    with open(json_dir, 'r') as fp:
        src_dict = json.load(fp)
    
    return src_dict


def select_from_json(img2pt_file_list, **kwargs):

    select_list = os.path.join(save_dir, 
    f"{scheme}_o{threshold_harris_opt}s{threshold_harris_sar}.json")

    # change the key in img2pt to wkt of sen12 instead of file path
    change_img2pt(img2pt_file_list)

    # select patch pairs with points in both modals
    select_patch(
        img2pt_json=img2pt_file_list, 
        save_path=select_list, 
        pt_range=[64, 192, 64, 192]
    )


def crn_cmd(read_func=open_img_np, fn2wkt=None, wkt2fn=None, pt_range=[]):

    parser = ArgumentParser()
    parser.add_argument(
        '-d', '--dset_path', help='folder of dataset', default=""
    )
    parser.add_argument(
        '-s', '--save_dir', help='folder to save results json', default=""
    )
    parser.add_argument(
        '-l', '--list_path', help='path of list file in which searched images are', default=""
    )
    parser.add_argument(
        '-t', '--threshold', help='threshold of opt and sar when harris', nargs="*", type=float, default=None
    )
    parser.add_argument(
        '-i', '--img2pt', help='list of img2pt path', nargs="*", type=str, default=None
    )
    args = parser.parse_args()


    workspace=""
    suffix_opt, suffix_sar= "", ""
    dset_path = args.dset_path
    save_dir = args.save_dir
    list_path = args.list_path
    threshold_abs = {
        "OPT" : 0.05, 
        "SAR" : 0.05
    }
    schema = "HARRIS"
    pt_range = [64, 192, 64, 192] if pt_range == [] else pt_range


    if args.img2pt is None:
        img2pt_file_list = get_png_crn(
            workspace=workspace, dset_path=dset_path, save_dir=save_dir, 
            list_path=list_path, threshold_abs=threshold_abs, 
            suffix_opt=suffix_opt, suffix_sar=suffix_sar, 
            read_func=read_func
        )

        # change the key in img2pt to wkt of sen12 instead of file path
        change_img2pt(img2pt_file_list, fn2wkt=fn2wkt)
        
    else:
        img2pt_file_list = args.img2pt

    if not args.threshold is None:
        threshold = {
            "PS-RGB" : args.threshold[0], 
            "SAR-Intensity" : args.threshold[1]
        }
    else:
        threshold = {
            "PS-RGB" : threshold_abs["OPT"], 
            "SAR-Intensity" : threshold_abs["SAR"]
        }

    # save path of selected patch pairs
    select_list = os.path.join(
        save_dir, 
        "{}_o{}s{}.json".format(
            schema, threshold["PS-RGB"], threshold["SAR-Intensity"]
        )
    )

    # select patch pairs with points in both modals
    select_patch(
        img2pt_json=img2pt_file_list, save_path=select_list, 
        pt_range=pt_range, wkt2fn=wkt2fn, thresh=threshold
    )


# test select_patch
if __name__ == "__main__1":
    img2pt_json = [
        "E:/workspace/SOMatch/json/sen12_ol_pt/img2pt_OPT_HARRISol_200.json",
        "E:/workspace/SOMatch/json/sen12_ol_pt/img2pt_SAR_HARRISol_200.json"
    ]
    save_path = "E:/workspace/SOMatch/json/sen12_ol_pt/select_patch.json"
    pt_range = [64, 192, 64, 192]
    thresh = {
        "PS-RGB": 3, 
        "SAR-Intensity": 0.6
    }

    select_patch(img2pt_json=img2pt_json, save_path=save_path, 
    pt_range=pt_range, thresh=thresh)


# change the filename in img2pt file 
# generated by func get_crn
if __name__ == "__main__1":
    img2pt_json = [
        "E:/workspace/SOMatch/tmp/json/sen12_ol_pt/subset_PS-RGB_HARRIS_subset200.json",
        "E:/workspace/SOMatch/tmp/json/sen12_ol_pt/subset_SAR_HARRIS_subset200.json"
    ]
    change_img2pt(img2pt_json)


# get key point from image in list file 
# and select the patch pairs with key points in both modals
if __name__ == "__main__1":
    import sys
    workspace_path = "E:/workspace/SOMatch/"
    sys.path.append(workspace_path)

    from tools.search_feature_points import get_crn_without_tif
    from utils.preprocess import bilatera_blur

    scheme = 'HARRIS'
    threshold_harris_opt = 1
    threshold_harris_sar = 0.5

    suffix1 = f"opt{threshold_harris_opt}" 
    suffix2 = f"sar{threshold_harris_sar}"
    # 如果json文件里是相对路径这里就要写数据集路径，否则写成空字符串
    dset_path = "E:/datasets/sen1-2"
    save_dir = "E:/workspace/SOMatch/json/sen12_spring_harris" 
    list_path = "E:/workspace/SOMatch/json/sen12_list/ROIs1158_spring.json"
    # save path of selected patch pairs
    select_list = os.path.join(save_dir, 
    f"{scheme}_o{threshold_harris_opt}s{threshold_harris_sar}.json")

    with open(list_path) as fp:
        dset_dict = json.load(fp)

    if not os.path.exists(save_dir):
        os.mkdir(save_dir)

    img2pt_file_list = []
    for modal, crn_dict in dset_dict.items():
        # 给相对路径加数据集路径
        if dset_path != '':
            for ids in range(len(crn_dict)):
                crn_dict[ids] = os.path.join(dset_path, crn_dict[ids])

        if modal == 'OPT':
            img_pt_dict = get_crn_without_tif(crn_dict, 
                os.path.join(save_dir,
                f"img2pt_{modal}_{scheme}{suffix1}.json"),
                type_='RGB', scheme=scheme, modal='PS-RGB', 
                border_radius=64, threshold_abs=threshold_harris_opt)
            img2pt_file_list.append(os.path.join(save_dir,
                f"img2pt_{modal}_{scheme}{suffix1}.json"))
        if modal == 'SAR':
            img_pt_dict = get_crn_without_tif(crn_dict, 
                os.path.join(save_dir,
                f"img2pt_{modal}_{scheme}{suffix2}.json"),
                type_='SAR', scheme=scheme,
                border_radius=64, threshold_abs=threshold_harris_sar, 
                pproc_func=bilatera_blur)
            img2pt_file_list.append(os.path.join(save_dir,
                f"img2pt_{modal}_{scheme}{suffix2}.json"))

    # change the key in img2pt to wkt of sen12 instead of file path
    change_img2pt(img2pt_file_list)

    # select patch pairs with points in both modals
    select_patch(img2pt_json=img2pt_file_list, 
    save_path=select_list, 
    pt_range=[64, 192, 64, 192])


if __name__ == "__main__":

    # change the key in img2pt to wkt of sen12 instead of file path
    change_img2pt(img2pt_file_list)

    # select patch pairs with points in both modals
    select_patch(img2pt_json=img2pt_file_list, 
    save_path=select_list, 
    pt_range=[64, 192, 64, 192])





